{{!

  Copyright (c) Meta Platforms, Inc. and affiliates.

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.

}}{{!
Generates a top-level file to be imported in the user's service code. It provides
wrappers for each of the service handlers that the user is then able to extend.
}}
{{> common/auto_generated_py}}
import builtins


from abc import ABCMeta
import typing as _typing

import folly.iobuf as _fbthrift_iobuf

import apache.thrift.metadata.thrift_types as _fbthrift_metadata
{{#program:generate_immutable_types}}
from {{program:base_library_package}}.serializer import serialize_iobuf, deserialize, Protocol
{{/program:generate_immutable_types}}
{{#program:generate_mutable_types}}
import {{program:base_library_package}}.mutable_containers as _fbthrift_python_mutable_containers
from {{program:base_library_package}}.mutable_serializer import serialize_iobuf, deserialize, Protocol
{{/program:generate_mutable_types}}
from {{program:base_library_package}}.server import ServiceInterface, RpcKind, PythonUserException

import {{program:module_path}}.{{> types/types_import_path}} as {{program:module_mangle}}__{{> types/types_import_path}}
import {{program:module_path}}.thrift_metadata as {{program:module_mangle}}__thrift_metadata
{{#program:include_namespaces}}
{{#has_services?}}
import {{included_module_path}}.thrift_services
{{/has_services?}}
{{#has_types?}}
import {{included_module_path}}.{{> types/types_import_path}} as {{included_module_mangle}}__{{> types/types_import_path}}
{{/has_types?}}
{{/program:include_namespaces}}
{{#program:adapter_modules}}
import {{module_path}}
{{/program:adapter_modules}}
{{#program:adapter_type_hint_modules}}
import {{module_path}}
{{/program:adapter_type_hint_modules}}

{{#program:services}}
class {{service:name}}Interface(
  {{#service:extends}}{{#service:external_program?}}
    {{service:module_path}}.thrift_services.{{/service:external_program?}}{{service:name}}Interface,
    {{/service:extends}}
    {{^service:extends?}}
    ServiceInterface,
    {{/service:extends?}}
    metaclass=ABCMeta
):

    @staticmethod
    def service_name() -> bytes:
        return b"{{service:name}}"

    def getFunctionTable(self) -> _typing.Mapping[bytes, _typing.Callable[..., object]]:
        functionTable = {
        {{#service:supported_service_functions}}
            b"{{function:name}}": ({{> services/function_kind }}, self._fbthrift__handler_{{function:name}}),
        {{/service:supported_service_functions}}
        }
        return {**super().getFunctionTable(), **functionTable}

    @staticmethod
    def __get_thrift_name__() -> str:
        return "{{program:name}}.{{service:name}}"

    @staticmethod
    def __get_metadata__() -> _fbthrift_metadata.ThriftMetadata:
        return {{program:module_mangle}}__thrift_metadata.gen_metadata_service_{{service:name}}()

    @staticmethod
    def __get_metadata_service_response__() -> _fbthrift_metadata.ThriftServiceMetadataResponse:
        return {{program:module_mangle}}__thrift_metadata._fbthrift_metadata_service_response_{{service:name}}()


{{#service:supported_service_functions}}

    {{^function:stream?}}async {{/function:stream?}}{{#function:stream_has_first_response?}}async {{/function:stream_has_first_response?}}def {{function:name}}(
            self{{#function:args}},
            {{field:py_name}}: {{#field:type}}{{> types/pep484_type}}{{/field:type}}{{/function:args}}
        ) -> {{#function:return_type}}{{> services/service_func_return_type }}{{/function:return_type}}:
        raise NotImplementedError("async def {{function:name}} is not implemented")

    {{#function:stream?}}
    async def _fbthrift__stream_wrapper_{{function:name}}(self, stream_generator: _typing.AsyncIterator[{{#function:stream_elem_type}}{{> types/pep484_type}}{{/function:stream_elem_type}}], protocol: Protocol) -> _typing.AsyncIterator[_fbthrift_iobuf.IOBuf]:
        {{#if (array.empty? function:stream.exceptions)}}
        async for item in stream_generator:
            yield serialize_iobuf({{> types/function_second_return_type}}(success=item), protocol)
        {{#else}}
        try:
            async for item in stream_generator:
                yield serialize_iobuf({{> types/function_second_return_type}}(success=item), protocol)
        {{#function:stream_exceptions}}{{#field:type}}
        except {{type:module_mangle}}.{{#type:structured}}{{> structs/unadapted_name}}{{/type:structured}} as e:
            return_struct = {{> types/function_second_return_type}}({{> fields/exception_field_name }}=e)
            buf = serialize_iobuf(return_struct, protocol)
            exp = PythonUserException('{{type:py3_namespace}}{{#type:structured}}{{> structs/unadapted_name}}{{/type:structured}}', str(e), buf)
            raise exp
        {{/field:type}}
        {{/function:stream_exceptions}}
        {{/if (array.empty? function:stream.exceptions)}}

    {{/function:stream?}}
    {{#function:sink?}}
    def _fbthrift__sink_wrapper_{{function:name}}(
        self,
        handler_sink_callback: {{> services/sink_callback }},
        protocol: Protocol
    ) -> _typing.Callable[[_typing.AsyncGenerator[_fbthrift_iobuf.IOBuf, None]], _typing.Awaitable[_fbthrift_iobuf.IOBuf]]:
        async def _fbthrift_iobuf_sink_callback(
            _iobuf_agen: _typing.AsyncGenerator[_fbthrift_iobuf.IOBuf, None],
        ) -> _fbthrift_iobuf.IOBuf:
            async def _handler_agen() -> _typing.AsyncGenerator[{{#function:sink_elem_type}}{{> types/pep484_type}}{{/function:sink_elem_type}}, None]:
                async for iobuf_item in _iobuf_agen:
                    sink_elem = deserialize(
                        {{> services/function_struct_prefix}}_result_sink_elem,
                        iobuf_item,
                        protocol
                    )
                    {{#if (array.empty? function:sink.exceptions)}}
                    assert sink_elem.success is not None, "unexpected empty sink element"
                    yield sink_elem.success
                    {{#else}}
                    if sink_elem.success is not None:
                        yield sink_elem.success
                        continue
                    for _ex_fld_name, ex in sink_elem:
                        if ex is not None:
                            raise ex
                    assert sink_elem.success is not None, "unexpected empty sink element"
                    {{/if (array.empty? function:sink.exceptions)}}

            {{#if (array.empty? function:sink.final_response_exceptions)}}
            final_resp = await handler_sink_callback(_handler_agen())
            return_struct = {{> services/function_struct_prefix}}_result_sink_final(success=final_resp)
            {{#else}}
            try:
                final_resp = await handler_sink_callback(_handler_agen())
                return_struct = {{> services/function_struct_prefix}}_result_sink_final(success=final_resp)
            {{#function:sink_final_response_exceptions}}{{#field:type}}
            except {{type:module_mangle}}.{{#type:structured}}{{> structs/unadapted_name}}{{/type:structured}} as ex:
                return_struct = {{> services/function_struct_prefix}}_result_sink_final({{> fields/exception_field_name }}=ex)
            {{/field:type}}{{/function:sink_final_response_exceptions}}
            {{/if (array.empty? function:sink.final_response_exceptions)}}
            return serialize_iobuf(return_struct, protocol)

        return _fbthrift_iobuf_sink_callback


    {{/function:sink?}}
    async def _fbthrift__handler_{{function:name}}(self, args: _fbthrift_iobuf.IOBuf, protocol: Protocol) {{!
            }}-> {{> services/service_handler_return_type }}:
        args_struct = deserialize({{> types/function_args_type}}, args, protocol)
        {{! TODO: call the postWrite hook here }}
        {{#function:exceptions?}}
        try:
            {{> services/get_handler_result }}
        {{#function:exceptions}}{{#field:type}}
        except {{type:module_mangle}}.{{#type:structured}}{{> structs/unadapted_name}}{{/type:structured}} as e:
            return_struct = {{> types/function_return_type}}({{> fields/exception_field_name }}=e)
            buf = serialize_iobuf(return_struct, protocol)
            exp = PythonUserException('{{type:py3_namespace}}{{#type:structured}}{{> structs/unadapted_name}}{{/type:structured}}', str(e), buf)
            raise exp
        {{/field:type}}
        {{/function:exceptions}}
        {{/function:exceptions?}}
        {{^function:exceptions?}}
        {{> services/get_handler_result }}
        {{/function:exceptions?}}
        {{^function:oneway?}}
        {{^function:sink_or_stream?}}
        return serialize_iobuf(return_struct, protocol)
        {{/function:sink_or_stream?}}
        {{#function:stream?}}
        return (serialize_iobuf(return_struct, protocol), return_stream)
        {{/function:stream?}}
        {{#function:sink?}}
        return (serialize_iobuf(return_struct, protocol), return_sink_callback)
        {{/function:sink?}}
        {{/function:oneway?}}

{{/service:supported_service_functions}}
{{/program:services}}
